# ros carla坐标绘制到地图中
import cv2
import numpy as np
import math
import rospy
from nav_msgs.msg import Odometry
import time
from collections import deque
from typing import Union, List
import threading
import argparse


def get_RT(map):
    R, T = None, None
    if map == 'Town01':
        # Town01的变换矩阵(ok)
        R = np.array(
            [
                [1.608978933495218, -0.0006133637127350534],
                [0.002967129565265202, -1.616573483032402],
            ]
        )
        T = np.array([[21.96085734252182], [73.98491744019127]])
        err_xy = 5
        loacl_map_h_w = 140
    elif map == 'Town02':
        # Town02的变换矩阵(ok)
        R = np.array(
            [
                [2.853288729947334, -0.001954564367013564],
                [-0.01084418875167121, -2.902583603546258],
            ]
        )
        T = np.array([[75.19143327247443], [-271.2606785755725]])
        err_xy = 1.5
        loacl_map_h_w = 150
    elif map == 'Town03':
        # Town03的变换矩阵(ok very nice)
        R = np.array(
            [
                [1.42539846369879, -0.000613806695560062],
                [-0.006309288645396612, -1.42529645694144],
            ]
        )
        T = np.array([[266.8903605584413], [321.8950052342699]])
        # T = np.array([[265], [320]])
        err_xy = 5
        loacl_map_h_w = 150
    elif map == 'Town04':
        # Town04的变换矩阵(ok)
        R = np.array(
            [
                [0.6945775929658323, -0.00242031540532944],
                [-0.001151290440462326, -0.7332215587816061],
            ]
        )
        T = np.array([[374.2232704087302], [303.5616340310737]])
        err_xy = 5
        loacl_map_h_w = 150
    elif map == 'Town05':
        # Town05的变换矩阵(not ok 高架桥误差大)
        R = np.array(
            [
                [1.151724847113483, 0.04979182850230186],
                [-0.04002015696056613, -1.193283702860338],
            ]
        )
        T = np.array([[387.2353315551624], [325.3770351204567]])
        err_xy = 5
        loacl_map_h_w = 128
    elif map == 'Town06':
        # Town06的变换矩阵（ok）
        R = np.array(
            [
                [0.622355330577449, 0.005832891756840823],
                [-0.002292540060410386, -0.5624124260904821],
            ]
        )
        T = np.array([[242.799258978577], [278.7354311222649]])
        err_xy = 5
        loacl_map_h_w = 128
    elif map == 'Town07':
        # Town07的变换矩阵(not ok 误差大)
        R = np.array(
            [
                [1.641298667149452, -0.01306631310848138],
                [0.03929811782735146, -1.530002461260377],
            ]
        )
        T = np.array([[461.2455550308022], [399.1818917817909]])
        err_xy = 1
        loacl_map_h_w = 128
    return R, T


def transform_img2carla(img_loc, R, T):
    """图像坐标系 转 carla世界坐标系

    Args:
        img_loc (np shape=(2,1)): 图像坐标系

    Returns:
        loc (np shape=(2,1)): carla世界坐标系

    """
    # carla = R^(-1)*(img-T)
    R_inv = np.linalg.inv(R)
    carla_loc = R_inv.dot(img_loc - T)
    return carla_loc


def transform_carla2img(loc, R, T):
    """carla世界坐标系 转 图像坐标系

    Args::
    Returns:
        img_loc (np shape=(2,1)): 图像坐标系
    """
    carla_loc = loc
    # img=R*carla+T
    img_loc = R.dot(carla_loc) + T
    # img_loc = np.matmul(R, carla_loc) + T
    # img_loc = carla_loc + T
    u = math.floor(max(img_loc[0, 0], 0))
    v = math.floor(max(img_loc[1, 0], 0))
    img_loc = np.array([[u], [v]])
    return img_loc


def draw_map(map, loc):
    pix_x = loc[0][0]
    pix_y = loc[1][0]

    # cv2.resizeWindow("MAP", 800, 600)
    cv2.circle(map, (pix_x, pix_y), 5, (255, 0, 0), -1)  # BRG 蓝色-起点
    # cv2.resizeWindow("MAP", 1920, 1080)
    cv2.imshow("MAP", map)
    cv2.waitKey(1)


class RosNode:
    def __init__(self, node_name='navigation_node', debug_mode=False):
        self.node_name = node_name
        self.carla_odometry_topic = '/carla/ego_vehicle/odometry'
        self.is_odometry_received = False
        self.carla_odometry_msg = Odometry()
        self.msg_pool = {}
        self.debug_mode = debug_mode

    def ros_init(self):
        rospy.init_node(self.node_name, anonymous=True, disable_signals=True)
        self.pool_regist(['odometry'])  # 注册消息池
        self.listener()

    def listener(self):
        rospy.Subscriber(
            self.carla_odometry_topic,
            Odometry,
            callback=self.callback_odometry,
            queue_size=1,
        )

    def callback_odometry(self, data: Odometry):
        this_time = rospy.get_time() * 1000.0  # ms
        self.is_odometry_received = True
        self.carla_odometry_msg = data
        self.msg_pool_append("odometry", data, this_time)
        if self.debug_mode:
            print(f"get {type(data)} msg")

    def msg_pool_append(self, msg_name: str, msg_data, msg_time: float):
        """消息池存储函数

        Args:
            msg_name (str): _description_
            msg_data (_type_): 消息data
            msg_time (rostime): 消息对应的ros时间
        """
        msg_tuple = (msg_data, msg_time)
        try:
            self.msg_pool[msg_name].append(msg_tuple)
        except KeyError:
            self.msg_pool[msg_name].append(msg_tuple)

    def pool_regist(self, msg_name: Union[str, List[str]]):
        """消息池中字典注册

        Args:
            msg_name (Union[str, List[str]]): 消息名，或者包含消息名的列表
        """
        if isinstance(msg_name, list):
            for m in msg_name:
                self.pool_regist(m)
        elif isinstance(msg_name, str):
            self.msg_pool[msg_name] = deque(maxlen=5)


class myThread(threading.Thread):
    def __init__(self, threadId, name):
        threading.Thread.__init__(self)
        self.threadId = threadId
        self.name = name
        self.counter = 0

    def run(self):
        global map,map_
        print("开始线程:", self.name)
        while 1:
            # 需要运行的程序：
            cv2.namedWindow("MAP", 0)
            cv2.setMouseCallback('MAP', draw_)
            if my_node.is_odometry_received:
                # 从ros中取消息
                carla_msg = my_node.msg_pool['odometry'][-1][0]
                loc_x = carla_msg.pose.pose.position.x
                loc_y = carla_msg.pose.pose.position.y
                loc_carla = np.array([[loc_x], [loc_y]])
                # carla坐标转成像素坐标
                loc_img = transform_carla2img(loc_carla, R, T)
                draw_map(map, loc_img)
                k = cv2.waitKey(1) & 0xFF
                if k == ord('c'):  # 按住c清空
                    map = cv2.imread(
                        f'/home/wrc/study/code/carla_tutorial/carlaMap/CARLAmap/{map_}.jpg'
                    )

            time.sleep(1e-2)  # 睡眠10ms
        cv2.destroyAllWindows()

    def __del__(self):
        print(self.name, "线程结束！")


def draw_(event, x, y, flag, param):
    global map, loc_carla, R, T
    if event == cv2.EVENT_LBUTTONDOWN:
        """左键单击打标签"""
        xy = "%d,%d" % (x, y)
        print(f"u:{x},v:{y}\n")
        img_loc = np.array([[x], [y]])
        carla_loc = transform_img2carla(img_loc=img_loc, R=R, T=T)
        print(f"carla loc :{carla_loc}\n")
        cv2.circle(map, (x, y), 2, (0, 255, 0), thickness=-1)


def get_argp():
    argparser = argparse.ArgumentParser(description=__doc__)
    argparser.add_argument(
        '--map',
        metavar='map',
        default='Town02',
        help='which map',
    )
    return argparser.parse_args()


if __name__ == "__main__":
    args = get_argp()
    map_ = args.map
    R, T = get_RT(map=map_)
    loc_carla = np.zeros((2, 1))
    loc_img = transform_carla2img(loc_carla, R, T)
    white_map = False
    if white_map:
        map = np.zeros([600, 800, 3], np.uint8)
        map.fill(255)
    else:
        map = cv2.imread(
            f'/home/wrc/study/code/carla_tutorial/carlaMap/CARLAmap/{map_}.jpg')
    # map_flip = cv2.flip(map, 1)
    print("按住c清除路径；左键点击该点对应carla坐标")
    # 2-ros主线程部分
    my_node = RosNode('carla_listen_node_1', debug_mode=False)
    my_node.ros_init()
    thread = myThread(1, 'myThread')
    thread.daemon = True
    thread.start()
    rospy.spin()
